# --------------------------------------------------------
# Relation Networks for Object Detection
# Copyright (c) 2017 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Modified by Jiayuan Gu, Dazhi Cheng, Yuwen Xiong
# --------------------------------------------------------
# Based on:
# MX-RCNN
# Copyright (c) 2016 by Contributors
# Licence under The Apache 2.0 License
# https://github.com/ijkguo/mx-rcnn/
# --------------------------------------------------------

import _init_paths
import argparse
import os
import sys
import time
import logging
from config.config import config, update_config


def parse_args():
    parser = argparse.ArgumentParser(description='Test a Faster R-CNN network')
    # general
    parser.add_argument('--cfg', help='experiment configure file name', required=True, type=str)

    args, rest = parser.parse_known_args()
    update_config(args.cfg)

    # rcnn
    parser.add_argument('--vis', help='turn on visualization', action='store_true')
    parser.add_argument('--ignore_cache', help='ignore cached results boxes', action='store_true')
    parser.add_argument('--thresh', help='valid detection threshold', default=1e-3, type=float)
    parser.add_argument('--shuffle', help='shuffle data on visualization', action='store_true')
    parser.add_argument('--test_epoch', help='the epoch model to be test', default=config.TEST.test_epoch, type=int)
    # nms
    parser.add_argument('--nms', help='params for nms or softnms', default=config.TEST.NMS, type=float)
    parser.add_argument('--softnms', help='whether to enable softnms', default=config.TEST.SOFTNMS, action='store_true')
    parser.add_argument('--naive_nms', help='whether to enable naive nms', default=False, action='store_true')
    parser.add_argument('--first_n', help='first_n for learn nms or nms', default=config.TEST.FIRST_N, type=int)
    parser.add_argument('--merge', help='merge method for learn nms', default=config.TEST.MERGE_METHOD, type=int)
    parser.add_argument('--debug', help='whether to enable debug mode', default=False, action='store_true')
    # dataset
    parser.add_argument('--test_set', help='which set to be tested', default=config.dataset.test_image_set, type=str)
    args, rest = parser.parse_known_args()
    # update config
    config.TEST.test_epoch = args.test_epoch
    config.TEST.NMS = args.nms
    config.TEST.SOFTNMS = args.softnms and (not args.naive_nms)
    config.TEST.FIRST_N = args.first_n
    config.TEST.MERGE_METHOD = args.merge
    config.dataset.test_image_set = args.test_set
    return args


args = parse_args()
curr_path = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, os.path.join(curr_path, '../external/mxnet', config.MXNET_VERSION))

import mxnet as mx
import numpy as np
from function.test_rcnn import test_rcnn
from utils.create_logger import create_logger


def main():
    ctx = [mx.gpu(int(i)) for i in config.gpus.split(',')]
    print args
    np.random.seed(0)
    mx.random.seed(0)
    logger, final_output_path = create_logger(config.output_path, args.cfg, config.dataset.test_image_set)

    test_rcnn(config, config.dataset.dataset, config.dataset.test_image_set, config.dataset.root_path, config.dataset.dataset_path,
              ctx, os.path.join(final_output_path, '..', '_'.join([iset for iset in config.dataset.image_set.split('+')]), config.TRAIN.model_prefix), config.TEST.test_epoch,
              args.vis, args.ignore_cache, args.shuffle, config.TEST.HAS_RPN, config.dataset.proposal, args.thresh, logger=logger, output_path=final_output_path)

if __name__ == '__main__':
    main()
